from ansible.module_utils.common.collections import ImmutableDict
from ansible.parsing.dataloader import DataLoader
from ansible.vars.manager import VariableManager
from ansible.inventory.manager import InventoryManager
from ansible.plugins.callback import CallbackBase
from ansible.executor.playbook_executor import PlaybookExecutor
from ansible import context
from ansible.utils.vars import load_extra_vars
from ansible.playbook.play import Play
from ansible.executor.task_queue_manager import TaskQueueManager
from django.utils import timezone


class ResultCallBack(CallbackBase):
    """
    重写callback
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        #以主机IP为key存放最后一步执行结果
        self.host_ok = {}
        self.host_failed = {}
        self.host_unreachable = {}
        self.host_skipped={}
        #所有主机执行信息列表
        self.host_all_info=[]
        self.host_last_stats={}
        self.host_stats=[]


    def v2_runner_on_ok(self, result, **kwargs):
        self.add_history(result,'ok')
        self.host_ok[result._host.get_name()] = result._result

    def v2_runner_on_failed(self, result, **kwargs):
        self.add_history(result,'failed')
        self.host_failed[result._host.get_name()] = result._result

    def v2_runner_on_unreachable(self, result, **kwargs):
        self.add_history(result,'unreachable')
        self.host_unreachable[result._host.get_name()] = result._result

    def v2_runner_on_skipped(self,result,**kwargs):
        self.add_history(result,'skipped')
        self.host_skipped[result._host.get_name()]=result._result

    def add_history(self,result,stats:str):
        result_dict=dict(
            exec_time=timezone.now(),
            host_ip=result._host.get_name(),
            task_name=result.task_name,
            result_info=result._result,
            stats=stats
        )
            
        self.host_all_info.append(result_dict)
        self.host_last_stats[result._host.get_name()]=result_dict



    
    def v2_playbook_on_stats(self,stats):
        hosts=sorted(stats.processed.keys())
        for h in hosts:
            t=stats.summarize(h)
            self.host_stats.append({
                h:t
            })
class AnsibleObj(object):
    def __init__(self, **kwargs):
        """
        初始化ansibel
        :param  connection:
        :param  module_path:
        :param  forks:
        :param  become:
        :param  become_method:
        :param  become_user:
        :param  check:
        :param  diff:
        :param  inventory:
        :param  variable_manager:
        """
        context.CLIARGS = ImmutableDict(
            connection=kwargs.get('connection', 'smart'),
            module_path=kwargs.get('module_path', None),
            forks=kwargs.get('forks', 10),
            remote_user=kwargs.get('remote_user', 'vagrant'),
            become=kwargs.get('become', None),
            become_method=kwargs.get('become_method', None),
            become_user=kwargs.get('become_user', None),
            check=kwargs.get('check', False),
            diff=kwargs.get('diff', False),
            verbosity=kwargs.get('verbosity', 5),
            syntax=None,
            start_at_task=None
        )
        self.dl = DataLoader()
        self.passwords = {
        }
        self.results_callback = ResultCallBack()
        self.inventory = InventoryManager(loader=self.dl,
                                          sources=kwargs.get('inventory', 'localhost'))
        self.variable_manager = kwargs.get('variable_manager',
                                           VariableManager(loader=self.dl, inventory=self.inventory))

    def run_command(self, command: str):
        """
        命令行控制台
        :param command:需要执行的命令
        """
        self.play_source = dict(
            name="Ansible Play",
            hosts='all',
            gather_facts='no',
            tasks=[
                dict(action=dict(module='shell', args=command),
                     register='shell_out'),
                dict(action=dict(module='debug', args=dict(
                    msg='{{shell_out.stdout}}')))
            ]
        )

        self.play = Play().load(self.play_source,
                                variable_manager=self.variable_manager, loader=self.dl)
        self.tqm = None
        try:
            self.tqm = TaskQueueManager(
                inventory=self.inventory,
                variable_manager=self.variable_manager,
                loader=self.dl,
                passwords=self.passwords,
                # Use our custom callback instead of the ``default`` callback plugin, which prints to stdout
                stdout_callback=self.results_callback
            )
            # most interesting data for a play is actually sent to the callback's methods
            self.result = self.tqm.run(self.play)
        except Exception as e:
            raise e
        finally:
            # we always need to cleanup child procs and the structures we use to communicate with them
            if self.tqm is not None:
                self.tqm.cleanup()


    def run_playbook(self, playbooks:str):
        self.playbook = PlaybookExecutor(playbooks=playbooks, inventory=self.inventory,
                                         variable_manager=self.variable_manager, loader=self.dl, passwords=self.passwords)
        try:
            self.playbook._tqm._stdout_callback=self.results_callback
            self.playbook.run()
        except Exception as e:
            raise e

    def run_script(self, scritp_path:str):
        """
        脚本控制台
        :param command:需要执行的命令
        """
        self.play_source = dict(
            name="Ansible Play",
            hosts='all',
            gather_facts='no',
            tasks=[
                dict(action=dict(module='script', args=scritp_path),
                     register='shell_out'),
                dict(action=dict(module='debug', args=dict(
                    msg='{{shell_out.stdout}}')))
            ]
        )

        self.play = Play().load(self.play_source,
                                variable_manager=self.variable_manager, loader=self.dl)
        self.tqm = None
        try:
            self.tqm = TaskQueueManager(
                inventory=self.inventory,
                variable_manager=self.variable_manager,
                loader=self.dl,
                passwords=self.passwords,
                # Use our custom callback instead of the ``default`` callback plugin, which prints to stdout
                stdout_callback=self.results_callback
            )
            # most interesting data for a play is actually sent to the callback's methods
            self.result = self.tqm.run(self.play)
        except Exception as e:
            raise e
        finally:
            # we always need to cleanup child procs and the structures we use to communicate with them
            if self.tqm is not None:
                self.tqm.cleanup()



    def get_result(self):
        result = dict(
            success=self.results_callback.host_ok,
            failed=self.results_callback.host_failed,
            unreachable=self.results_callback.host_unreachable,
            skipped=self.results_callback.host_skipped,
            history=self.results_callback.host_all_info,
            stats=self.results_callback.host_stats,
            last_stats=self.results_callback.host_last_stats
        )
        return result


if __name__ == "__main__":
    inventory = '/tmp/tmp_inventory/'
    ar = AnsibleObj(inventory=inventory)
    ar.run_command('pwd')
    #ar.run_playbook(['/home/bolinz/下载/slb-ansible-dev/main.yml'])
    #ar.run_script('/home/bolinz/123')
    result = ar.get_result()
    print(result)
